#!/usr/bin/env bash
#
# This script runs inside chrooted environment. It installs grub and its
# Configuration file.
#

set -o errexit
set -o pipefail
set -o xtrace

export DEBIAN_FRONTEND=noninteractive

export APT_OPTIONS="-oAPT::Install-Recommends=false \
		  -oAPT::Install-Suggests=false \
		    -oAcquire::Languages=none"

if [ $(dpkg --print-architecture) = "amd64" ];
then
	ARCH="amd64";
else
	ARCH="arm64";
fi

# Mirror fallback function for resilient apt-get update
function apt_update_with_fallback {
	local sources_file="/etc/apt/sources.list"
	local max_attempts=2
	local attempt=1

	# Detect the current region from sources.list (it's already been substituted)
	# Extract the region from existing sources.list entries
	local current_region=$(grep -oP '(?<=http://)[^.]+(?=\.clouds\.ports\.ubuntu\.com)' "${sources_file}" | head -1 || echo "")

	# Define mirror tiers (in priority order)
	local -a mirror_tiers=(
		"${current_region}.clouds.ports.ubuntu.com"  # Tier 1: Regional CDN (as set in sources.list)
		"ports.ubuntu.com"                             # Tier 2: Global pool
	)

	# If we couldn't detect current region, skip tier 1
	if [ -z "${current_region}" ]; then
		echo "Warning: Could not determine region from sources.list, skipping regional CDN"
		mirror_tiers=("${mirror_tiers[@]:1}")  # Remove first element
	fi

	for mirror in "${mirror_tiers[@]}"; do
		echo "========================================="
		echo "Attempting apt-get update with mirror: ${mirror}"
		echo "Attempt ${attempt} of ${max_attempts}"
		echo "========================================="

		# Update sources.list to use current mirror
		sed -i "s|http://[^/]*/ubuntu-ports/|http://${mirror}/ubuntu-ports/|g" "${sources_file}"

		# Show what we're using
		echo "Current sources.list configuration:"
		grep -E '^deb ' "${sources_file}" | head -3

		# Attempt update with timeout (5 minutes)
		if timeout 300 apt-get $APT_OPTIONS update 2>&1; then
			echo "========================================="
			echo "✓ Successfully updated apt cache using mirror: ${mirror}"
			echo "========================================="
			return 0
		else
			local exit_code=$?
			echo "========================================="
			echo "✗ Failed to update using mirror: ${mirror}"
			echo "Exit code: ${exit_code}"
			echo "========================================="

			# Clean partial downloads
			apt-get clean
			rm -rf /var/lib/apt/lists/*

			# Exponential backoff before next attempt
			if [ ${attempt} -lt ${max_attempts} ]; then
				local sleep_time=$((attempt * 5))
				echo "Waiting ${sleep_time} seconds before trying next mirror..."
				sleep ${sleep_time}
			fi
		fi

		attempt=$((attempt + 1))
	done

	echo "========================================="
	echo "ERROR: All mirror tiers failed after ${max_attempts} attempts"
	echo "========================================="
	return 1
}



function update_install_packages {
	source /etc/os-release

	# Update APT with new sources (using fallback mechanism)
	cat /etc/apt/sources.list
	if ! apt_update_with_fallback; then
		echo "FATAL: Failed to update package lists with any mirror tier"
		exit 1
	fi
	apt-get $APT_OPTIONS --yes dist-upgrade

	# Do not configure grub during package install
	if [ "${ARCH}" = "amd64" ]; then
		echo 'grub-pc grub-pc/install_devices_empty select true' | debconf-set-selections
		echo 'grub-pc grub-pc/install_devices select' | debconf-set-selections
	# Install various packages needed for a booting system
		apt-get install -y \
		linux-aws \
		grub-pc \
		e2fsprogs
	else
		apt-get install -y e2fsprogs
	fi
	# Install standard packages
	apt-get install -y \
		sudo \
		wget \
		cloud-init \
		acpid \
		ec2-hibinit-agent \
		ec2-instance-connect \
		hibagent \
		ncurses-term \
		ssh-import-id \

	# apt upgrade
	apt-get upgrade -y

	# Install OpenSSH and other packages
	sudo add-apt-repository --yes universe
	if ! apt_update_with_fallback; then
		echo "FATAL: Failed to update package lists after adding universe repository"
		exit 1
	fi
	apt-get install -y --no-install-recommends \
		openssh-server \
		git \
		ufw \
		cron \
		logrotate \
		fail2ban \
		locales \
		at \
		less \
		python3-systemd

	if [ "${ARCH}" = "arm64" ]; then
		apt-get $APT_OPTIONS --yes install linux-aws initramfs-tools dosfstools
	fi
}

function setup_locale {
cat << EOF >> /etc/locale.gen
en_US.UTF-8 UTF-8
EOF

cat << EOF > /etc/default/locale
LANG="C.UTF-8"
LC_CTYPE="C.UTF-8"
EOF
	locale-gen en_US.UTF-8
}

function setup_postgesql_env {
	    # Create the directory if it doesn't exist
    sudo mkdir -p /etc/environment.d
    
    # Define the contents of the PostgreSQL environment file
    cat <<EOF | sudo tee /etc/environment.d/postgresql.env >/dev/null
LOCALE_ARCHIVE=/usr/lib/locale/locale-archive
LANG="en_US.UTF-8"
LANGUAGE="en_US.UTF-8"
LC_ALL="en_US.UTF-8"
LC_CTYPE="en_US.UTF-8"
EOF
}

function install_packages_for_build {
	apt-get install -y --no-install-recommends linux-libc-dev \
	 acl \
	 magic-wormhole sysstat \
	 build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libsystemd-dev libpq-dev libxml2-utils uuid-dev xsltproc ssl-cert \
	 gcc-10 g++-10 \
	 libgeos-dev libproj-dev libgdal-dev libjson-c-dev libboost-all-dev libcgal-dev libmpfr-dev libgmp-dev cmake \
	 libkrb5-dev \
	 maven default-jre default-jdk \
	 curl gpp apt-transport-https cmake libc++-dev libc++abi-dev libc++1 libglib2.0-dev libtinfo5 libc++abi1 ninja-build python \
	 liblzo2-dev

	source /etc/os-release

	apt-get install -y --no-install-recommends llvm-11-dev clang-11
	# Mark llvm as manual to prevent auto removal
	apt-mark manual libllvm11:arm64
}

function setup_apparmor {
	apt-get install -y apparmor apparmor-utils auditd

	# Copy apparmor profiles
	cp -rv /tmp/apparmor_profiles/* /etc/apparmor.d/
}

function setup_grub_conf_arm64 {
cat << EOF > /etc/default/grub
GRUB_DEFAULT=0
GRUB_TIMEOUT=0
GRUB_TIMEOUT_STYLE="hidden"
GRUB_DISTRIBUTOR="Supabase postgresql"
GRUB_CMDLINE_LINUX_DEFAULT="nomodeset console=tty1 console=ttyS0 ipv6.disable=0"
EOF
}

# Install GRUB
function install_configure_grub {
	if [ "${ARCH}" = "arm64" ]; then
		apt-get $APT_OPTIONS --yes install cloud-guest-utils fdisk grub-efi-arm64 efibootmgr
		setup_grub_conf_arm64
		rm -rf /etc/grub.d/30_os-prober
		sleep 1
	fi
	grub-install /dev/xvdf && update-grub
}

# skip fsck for first boot
function disable_fsck {
	touch /fastboot
}

# Don't request hostname during boot but set hostname
function setup_hostname {
	# Set the static hostname
	echo "ubuntu" > /etc/hostname
	chmod 644 /etc/hostname
	# Update netplan configuration to not send hostname
	cat << EOF > /etc/netplan/01-hostname.yaml
network:
  version: 2
  ethernets:
    eth0:
      dhcp4: true
      dhcp4-overrides:
        send-hostname: false
EOF
	# Set proper permissions for netplan security
	chmod 600 /etc/netplan/01-hostname.yaml
}

# Set options for the default interface
function setup_eth0_interface {
cat << EOF > /etc/netplan/eth0.yaml
network:
  version: 2
  ethernets:
    eth0:
      dhcp4: true
EOF
	# Set proper permissions for netplan security
	chmod 600 /etc/netplan/eth0.yaml
}

function disable_sshd_passwd_auth {
	sed -i -E -e 's/^#?\s*PasswordAuthentication\s+(yes|no)\s*$/PasswordAuthentication no/g' \
	  -e 's/^#?\s*ChallengeResponseAuthentication\s+(yes|no)\s*$/ChallengeResponseAuthentication no/g' \
	 /etc/ssh/sshd_config
}

function create_admin_account {
	groupadd admin
}

#Set default target as multi-user
function set_default_target {
	rm -f /etc/systemd/system/default.target
	ln -s /lib/systemd/system/multi-user.target /etc/systemd/system/default.target
}

# Setup ccache
function setup_ccache {
	apt-get install ccache -y
	mkdir -p /tmp/ccache
	export PATH=/usr/lib/ccache:$PATH
	echo "PATH=$PATH" >> /etc/environment
}

# Clear apt caches
function cleanup_cache {
	apt-get clean
}

update_install_packages
setup_locale
setup_postgesql_env
#install_packages_for_build
install_configure_grub
setup_apparmor
setup_hostname
create_admin_account
set_default_target
setup_eth0_interface
disable_sshd_passwd_auth
disable_fsck
#setup_ccache
cleanup_cache
